# 波束形成优化
import argparse
import sys
from typing import Dict
import matplotlib.pyplot as plt
import numpy as np
import torch
from torch.optim.lr_scheduler import ReduceLROnPlateau
from scipy.signal import windows as SSW # import chebwin
from apps.wfs.bfo_config import BfoConfig as BG
from rsp.antenna.array import linear_array as LA

class BfoApp(object):
    def __init__(self):
        self.name = 'apps.wfs.bfo_app.BfoApp'

    @staticmethod
    def startup(params:Dict = {}) -> None:
        print(f'波束成形（复数）优化 v0.0.4')
        # BfoApp.bf_scan_demo(params=params)
        BfoApp.draw_bf_main(params=params)
        exit(0)
        # BfoApp.get_tc()
        # # 准备数据集
        # params['ds_mode'] = 0 # 生成32个天线的标签集
        # BfoApp.generate_ds(params=params)
        # params['ds_mode'] = 1 # 生成8个天线的数据集
        # BfoApp.generate_ds(params=params)
        # 训练网络
        # BfoApp.train_main(params=params)
        # BfoApp.compare_training_result(params=params)

        
        
        w = []
        w_rank = 4
        for i in range(BG.N):
            coefs = torch.randn(w_rank)
            w_i = lambda x: sum(coef * x**i for i, coef in enumerate(coefs))
            w.append(w_i)
        real = torch.randn(BG.N, dtype=torch.float64)
        imag = torch.randn(BG.N, dtype=torch.float64)
        a = torch.complex(real, imag)
        print(f'a: {a.shape}; {a.dtype};')
        v1 = torch.sum(torch.tensor([f_i(a[i]) for i, f_i in enumerate(w)]))
        print(f'### v1: {v1}; a: {type(a)}; {a.dtype};')

    @staticmethod
    def train_main(params:Dict = {}) -> None:
        '''
        from Tschebyscheff loss=0.09756
        '''
        # 载入波束成形相关内容和训练目标
        X = torch.load('./work/wfs/X.pt', weights_only=True) # 导向向量
        X = X.to(BG.device)
        y = torch.load('./work/wfs/y_wa.pt', weights_only=True).unsqueeze(1)
        y = y.to(BG.device)
        print(f'X: {X.shape}, {X.dtype}; y: {y.shape}, {y.dtype};')
        af = SSW.chebwin(BG.N, at=BG.SSL_db, sym=True)
        phases = np.exp(1j * 2 * np.pi / BG.lambda_ * np.arange(BG.N) * BG.d[0] * np.sin(BG.theta/180.0*np.pi))
        w_np = af * phases # 权重向量初始值为8天线切比雪夫权重
        # 生成随机权重
        # w = torch.from_numpy(w_np) # 取切比雪夫初始值
        real_part = torch.randn((BG.N,), dtype=torch.float64)  # 生成实部的随机数Tensor
        imag_part = torch.randn((BG.N,), dtype=torch.float64) 
        w = torch.complex(real_part, imag_part)
        w.requires_grad = True
        w = w.to(BG.device)
        # loss_fn = torch.nn.MSELoss()
        loss_fn = torch.nn.L1Loss()
        mparams = [{'params': torch.nn.Parameter(w)}] # 模型需要学习的参数
        # 预热阶段
        print(f'##### 开始预热训练')
        warmup_opt = torch.optim.AdamW(mparams, lr=1e-6, betas=(0.9, 0.99), weight_decay=1e-2)
        warmup_scheduler = ReduceLROnPlateau(warmup_opt, 'min', threshold=1e-4, eps=1e-8)
        warmup_epochs = 5
        for epoch in range(warmup_epochs):
            BfoApp.train_epoch(epoch, mparams, loss_fn, warmup_opt, warmup_scheduler, X, y)
        # 正式训练过程
        print(f'##### 开始正式训练')
        optimizer = torch.optim.AdamW(mparams, lr=1e-5, betas=(0.9, 0.99), weight_decay=1e-2)
        scheduler = ReduceLROnPlateau(optimizer, 'min', threshold=1e-4, eps=1e-8)
        epochs = 2000000
        best_loss = sys.float_info.max
        improve_threshold = 0.0000001
        cumulative_steps = 0
        max_unimproves = 500
        pt_fn = './work/wfs/wfs.pt'
        for epoch in range(epochs):
            loss = BfoApp.train_epoch(epoch, mparams, loss_fn, optimizer, scheduler, X, y)
            print(f'epoch_{epoch}: loss={loss};')
            cumulative_steps += 1
            if best_loss > loss:
                if best_loss - loss > improve_threshold:
                    torch.save(mparams, pt_fn)
                    cumulative_steps = 0
                    best_loss = loss
            if cumulative_steps > max_unimproves:
                # y = model(X)
                print(f'Earlly Stopping!!!!!!')
                break
        # 训练后的结果评估
        a_np = X.cpu().numpy()
        # 天线图响应
        B = np.dot(a_np.transpose(), np.conj(w_np))
        B = np.abs(B) / np.max(np.abs(B))
        # 绘制波束形成图
        fig, ax = plt.subplots(3, 1, sharex=True, sharey=True, figsize=(7, 2), dpi=300)
        plt.subplots_adjust(wspace=0.1, hspace=0.05)
        ax[0].plot(BG.thetas[0], 20 * np.log10(B), '--')
        ax[0].set_title('Initial', fontsize=12, color='blue', loc='center')
        #
        w_hat_np = mparams[0]["params"][0].detach().cpu().numpy()
        B1 = np.dot(a_np.transpose(), np.conj(w_hat_np))
        B1 = np.abs(B1) / np.max(np.abs(B1))
        # 绘制波束形成图
        ax[1].plot(BG.thetas[0], 20 * np.log10(B1), '--')
        ax[1].set_title('Learned', fontsize=12, color='blue', loc='center')
        #
        y_np = y.detach().cpu().numpy()
        ax[2].plot(BG.thetas[0], 20.0 * np.log10(abs(y_np) + np.finfo(float).eps), '')
        ax[2].set_title('dest', fontsize=12, color='blue', loc='center')
        ax[2].set_ylim(-120.0, 0.5)
        plt.show()
        print(f'^_^ The End! ^_^')

    @staticmethod
    def compare_training_result(params:Dict = {}) -> None:
        # 载入波束成形相关内容和训练目标
        X = torch.load('./work/wfs/X.pt', weights_only=True) # 导向向量
        X = X.to(BG.device)
        y = torch.load('./work/wfs/y_wa.pt', weights_only=True).unsqueeze(1)
        y = y.to(BG.device)
        print(f'X: {X.shape}, {X.dtype}; y: {y.shape}, {y.dtype};')
        af = SSW.chebwin(BG.N, at=BG.SSL_db, sym=True)
        phases = np.exp(1j * 2 * np.pi / BG.lambda_ * np.arange(BG.N) * BG.d[0] * np.sin(BG.theta/180.0*np.pi))
        w_np = af * phases # 权重向量初始值为8天线切比雪夫权重
        # 读入训练好的权重
        w = torch.load('./work/wfs/wfs.pt', weights_only=True)
        print(f'### w: {type(w[0]["params"][0])}; ????????????')
        w = w[0]["params"][0]
        w = w.to(BG.device)
        # 训练后的结果评估
        a_np = X.detach().cpu().numpy()
        # 天线图响应
        B = np.dot(a_np.transpose(), np.conj(w_np))
        B = np.abs(B) / np.max(np.abs(B))
        # 绘制波束形成图
        fig, ax = plt.subplots(3, 1, sharex=True, sharey=True, figsize=(7, 2), dpi=300)
        plt.subplots_adjust(wspace=0.1, hspace=0.05)
        ax[0].plot(BG.thetas[0], 20 * np.log10(B), '--')
        ax[0].set_title('Initial', fontsize=12, color='blue', loc='center')
        #
        w_hat_np = w.detach().cpu().numpy()
        B1 = np.dot(a_np.transpose(), np.conj(w_hat_np))
        B1 = np.abs(B1) / np.max(np.abs(B1))
        # 绘制波束形成图
        ax[1].plot(BG.thetas[0], 20 * np.log10(B1), '--')
        ax[1].set_title('Learned', fontsize=12, color='blue', loc='center')
        #
        y_np = y.detach().cpu().numpy()
        ax[2].plot(BG.thetas[0], 20.0 * np.log10(abs(y_np) + np.finfo(float).eps), '')
        ax[2].set_title('dest', fontsize=12, color='blue', loc='center')
        ax[2].set_ylim(-120.0, 0.5)
        plt.show()
        print(f'^_^ The End! ^_^')
            

    @staticmethod
    def train_epoch(epoch, model_params, loss_fn, optimizer, scheduler, X, y_gt):
        batch_size = 100
        total_loss = torch.tensor([0.0]).to(X.device)
        for bi in range(batch_size):
            # train
            y_hat = torch.matmul(X.T, torch.conj(model_params[0]['params'][0]).unsqueeze(1))
            y_hat = torch.abs(y_hat) / torch.max(torch.abs(y_hat))
            loss = loss_fn(y_gt, y_hat)
            total_loss += loss
            optimizer.zero_grad()
            loss.backward()
            for param in model_params[0]['params']:
                if param.grad is not None:
                    param.grad = param.grad.resolve_conj()
            optimizer.step()
        if scheduler is not None:
            scheduler.step(total_loss)
        return total_loss.detach().cpu().item()/batch_size




        

    @staticmethod
    def generate_ds(params:Dict = {}) -> None:
        old_N = BG.N
        if params['ds_mode'] == 1:
            BG.N = 8
        else:
            BG.N = 32
        af = SSW.chebwin(BG.N, at=BG.SSL_db, sym=True)
        phases = np.exp(1j * 2 * np.pi / BG.lambda_ * np.arange(BG.N) * BG.d[0] * np.sin(BG.theta/180.0*np.pi))
        w_vec = af * phases # 权重向量
        # 导向向量
        a_vec = np.array([np.exp(1j * 2 * np.pi / BG.lambda_ * np.arange(BG.N) * BG.d[0] * np.sin((ti))) for ti in BG.thetas_rad[0]]).T
        # 天线图响应
        B = np.dot(a_vec.transpose(), np.conj(w_vec))
        B = np.abs(B) / np.max(np.abs(B))
        print(f'w_vec: {w_vec.shape}; a_vec: {a_vec.shape}; B: {B.shape}')
        if params['ds_mode'] == 1:
            print(f'生成样本集（8天线阵列）')
            X = torch.from_numpy(a_vec)
            torch.save(X, './work/wfs/X.pt')
        else:
            print(f'生成标签集（32天线阵列）')
            y_wa = torch.from_numpy(B)
            torch.save(y_wa, './work/wfs/y_wa.pt')
        BG.N = old_N

    @staticmethod
    def get_tc() -> None:
        print(f'获取切比雪夫系数并转化为权重向量和导向向量')
        af = SSW.chebwin(BG.N, at=BG.SSL_db, sym=True)
        offset = int(np.floor(BG.N / 2))
        coefficients = np.roll(af, offset) # 把后一半倒到前面
        # coefficients[i] * cos((i + 0.5) * psi) for i in range(offset): (1000, 16)
        # Calculate the wavenumber
        k = 2.0 * np.pi * BG.f / BG.c
        # Calculate the phase
        psi = k * BG.d[0] * (np.cos(np.radians(BG.thetas[0])) - np.cos(BG.theta/180.0*np.pi))
        afx = sum(coefficients[i] * np.cos((i + 0.5) * psi) for i in range(offset)) # af: (1000,)
        afx = afx / np.amax(abs(afx)) # 对其进行归一化
        phases = np.exp(1j * 2 * np.pi / BG.lambda_ * np.arange(BG.N) * BG.d[0] * np.sin(BG.theta/180.0*np.pi))
        print(f'pahses: {phases};')
        w_vec = af * phases
        a_vec = np.array([np.exp(1j * 2 * np.pi / BG.lambda_ * np.arange(BG.N) * BG.d[0] * np.sin((ti))) for ti in BG.thetas_rad[0]]).T
        B = np.dot(a_vec.transpose(), np.conj(w_vec))
        B = np.abs(B) / np.max(np.abs(B))
        print(f'w_vec: {w_vec.shape}; a_vec: {a_vec.shape}; B: {B.shape}')
        # 绘制波束形成图
        fig, ax = plt.subplots(1, 2, sharex=True, sharey=True, figsize=(7, 2), dpi=300)
        plt.subplots_adjust(wspace=0.1, hspace=0.05)
        ax[0].plot(BG.thetas[0], 20 * np.log10(B), '--')
        ax[0].set_title('w dot a(Compex number)', fontsize=12, color='blue', loc='center')
        #
        ax[1].plot(BG.thetas[0], 20.0 * np.log10(abs(afx) + np.finfo(float).eps), '')
        ax[1].set_title('AF(Real number)', fontsize=12, color='blue', loc='center')
        ax[1].set_ylim(-120.0, 0.5)
        plt.show()
        print(f'^_^ The End! ^_^')


    @staticmethod
    def draw_bf_main(params:Dict = {}) -> None:
        theta_0 = 30 / 180.0 * np.pi
        w = BfoApp.cal_tschebyscheff_w(BG.N, BG.lambda_, BG.d[0], theta_0, BG.SSL_db)
        a = np.exp(1j * 2 * np.pi * np.sin(BG.thetas_rad) * BG.n * BG.d / BG.lambda_)
        print(f'### {w};')
        BfoApp.draw_bf(w, a)
    @staticmethod
    def draw_bf(w, a) -> None:
        B = np.dot(a.transpose(), np.conj(w))
        B = np.abs(B) / np.max(np.abs(B))
        # 绘制波束形成图
        fig, ax = plt.subplots(1, 2, sharex=True, sharey=True, figsize=(7, 2), dpi=300)
        plt.subplots_adjust(wspace=0.1, hspace=0.05)
        ax[0].plot(BG.thetas.squeeze(0), 20 * np.log10(B), '--')
        plt.show()

    @staticmethod
    def cal_tschebyscheff_w(N, lambda_, d, theta_0, ssl) -> np.ndarray:
        # w = tschebyscheff_weights(N, wavelength, d, theta_0, alpha_linear)
        # 计算x值
        x = np.cos(np.pi / N * np.arccos(ssl))
        # 初始化切比雪夫多项式
        T = np.zeros(N)
        T[0] = 1
        T[1] = x
        # 计算切比雪夫多项式
        for n in range(2, N):
            T[n] = 2 * x * T[n - 1] - T[n - 2]
        # 计算权重向量
        weights = np.exp(1j * 2 * np.pi / lambda_ * d * np.arange(N) * np.sin(theta_0)) * T
        # 归一化权重向量
        weights /= np.dot(np.conjugate(np.exp(1j * 2 * np.pi / lambda_ * d * np.arange(N) * np.sin(theta_0))), weights)
        return weights


    @staticmethod
    def bf_scan_demo(params:Dict = {}) -> None:
        '''
        波束形成图绘制（给定导向向量），以及波束扫描方位谱
        '''
        a = np.exp(1j * 2 * np.pi * np.sin(BG.thetas_rad) * BG.n * BG.d / BG.lambda_)
        w = np.exp(1j * 2 * np.pi * np.sin(BG.theta_rad) * BG.n * BG.d / BG.lambda_)
        B = np.dot(a.transpose(), np.conj(w))
        B = np.abs(B) / np.max(np.abs(B))
        B = B.squeeze(1)
        print(f'????? B: {B.shape};')
        # 绘制波束形成图
        fig, ax = plt.subplots(1, 2, sharex=True, sharey=True, figsize=(7, 2), dpi=300)
        plt.subplots_adjust(wspace=0.1, hspace=0.05)
        ax[0].plot(BG.thetas.squeeze(0), 20 * np.log10(B), '--')
        a = np.mat(np.exp(1j * 2 * np.pi * np.sin(BG.theta_rad) * BG.n * BG.d / BG.lambda_))
        w = np.mat(np.exp(1j * 2 * np.pi * np.sin(BG.thetas_rad) * BG.n * BG.d / BG.lambda_))
        sigma = np.transpose(np.conj(w)) * a * np.transpose(np.conj(a)) * w
        sigma = np.diag(sigma)
        sigma = sigma / np.max(sigma)
        ax[1].plot(BG.thetas.squeeze(0), 10 * np.log10(sigma))
        plt.show()

def main(params:Dict = {}) -> None:
    BfoApp.startup(params=params)

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--run_mode', action='store',
        type=int, default=1, dest='run_mode',
        help='run mode'
    )
    return parser.parse_args()

if '__main__' == __name__:
    args = parse_args()
    params = vars(args)
    main(params=params)